# -*- coding: utf-8 -*-
r"""
Intel 场景分类 - HOG+颜色直方图 + 随机森林 + CV
输出文件统一加 small_ 前缀，模型保存到 D:\codes\wcytwo\results
"""
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from skimage.feature import hog
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
import joblib
import warnings
warnings.filterwarnings("ignore")

# =====================  CONFIG  =====================
CODE_FILE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_ROOT     = os.path.dirname(CODE_FILE_DIR)
TRAIN_DIR     = os.path.join(DATA_ROOT, "seg_train")
TEST_DIR      = os.path.join(DATA_ROOT, "seg_test")
CODE_DIR      = os.path.join(DATA_ROOT, "codes")
RESULT_DIR    = os.path.join(DATA_ROOT, "results")

IMG_SIZE        = (128, 128)
HOG_WEIGHT      = 0.7
COLOR_WEIGHT    = 0.3
ORIENTATIONS    = 9
PIXELS_PER_CELL = (8, 8)
CELLS_PER_BLOCK = (2, 2)

FAST_DEBUG    = True
MAX_PER_CLASS = 200
if FAST_DEBUG:
    RF_PARAM_GRID = {
        "n_estimators": [200],
        "max_depth": [None],
        "min_samples_split": [2],
        "min_samples_leaf": [1],
        "n_jobs": [1]
    }
    CV_FOLDS = 3
else:
    RF_PARAM_GRID = {
        "n_estimators": [200, 400],
        "max_depth": [None, 20],
        "min_samples_split": [2, 5],
        "min_samples_leaf": [1, 2],
        "n_jobs": [-1]
    }
    CV_FOLDS = 5

RANDOM_STATE = 42
CLASS_NAMES  = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
# ====================================================

os.makedirs(CODE_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

# 新增：模型保存目录
MODEL_SAVE_DIR = r"D:\codes\wcytwo\results"
os.makedirs(MODEL_SAVE_DIR, exist_ok=True)


def extract_feature(img_path):
    img = cv2.imread(img_path)
    if img is None:
        raise ValueError(f"无法读取 {img_path}")
    img = cv2.resize(img, IMG_SIZE)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    hog_feat = hog(gray, orientations=ORIENTATIONS,
                   pixels_per_cell=PIXELS_PER_CELL,
                   cells_per_block=CELLS_PER_BLOCK,
                   block_norm='L2-Hys', feature_vector=True)
    color_hist = []
    for ch in range(3):
        h = cv2.calcHist([img], [ch], None, [64], [0, 256])
        h = cv2.normalize(h, h).flatten()
        color_hist.append(h)
    color_feat = np.concatenate(color_hist)
    return np.hstack([HOG_WEIGHT * hog_feat, COLOR_WEIGHT * color_feat])


def load_dataset(folder_dir):
    X, y, paths = [], [], []
    class_to_idx = {name: i for i, name in enumerate(CLASS_NAMES)}
    for cls_name in CLASS_NAMES:
        cls_dir = os.path.join(folder_dir, cls_name)
        if not os.path.isdir(cls_dir):
            continue
        fnames = [f for f in os.listdir(cls_dir)
                  if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        if FAST_DEBUG:
            fnames = fnames[:MAX_PER_CLASS]
        for fname in tqdm(fnames, desc=cls_name):
            full_path = os.path.join(cls_dir, fname)
            X.append(extract_feature(full_path))
            y.append(class_to_idx[cls_name])
            paths.append(full_path)
    return np.array(X, dtype=np.float32), np.array(y, dtype=np.int32), paths


def train():
    print("加载训练集...")
    X, y, _ = load_dataset(TRAIN_DIR)

    print("随机森林 GridSearchCV 调参...")
    rf = RandomForestClassifier(random_state=RANDOM_STATE)
    n_jobs = RF_PARAM_GRID.get("n_jobs", [1])[0]
    grid = GridSearchCV(rf, RF_PARAM_GRID, cv=CV_FOLDS,
                        scoring='accuracy', n_jobs=n_jobs, verbose=2)
    grid.fit(X, y)

    print(f"最佳参数: {grid.best_params_}")
    print(f"最佳交叉验证 Accuracy: {grid.best_score_:.4f}")

    best_model = grid.best_estimator_
    # 改成 small_ 前缀并保存到指定目录
    model_path = os.path.join(MODEL_SAVE_DIR, "small_best_model.pkl")
    joblib.dump(best_model, model_path)
    print(f"模型已保存 -> {model_path}")
    return best_model


def evaluate(model):
    print("加载测试集...")
    X_test, y_test, paths = load_dataset(TEST_DIR)
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    print(f"测试集 Accuracy: {acc:.4f}")

    cm = confusion_matrix(y_test, y_pred, labels=range(len(CLASS_NAMES)))
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=CLASS_NAMES)
    disp.plot(cmap='Blues', xticks_rotation=45)
    plt.tight_layout()
    plt.savefig(os.path.join(RESULT_DIR, "small_confusion_matrix.png"), dpi=300)
    plt.close()

    show_predictions(paths, y_test, y_pred, 12)
    return acc


def show_predictions(paths, y_test, y_pred, n_show=12):
    import random
    idxs = random.sample(range(len(paths)), n_show)
    plt.figure(figsize=(12, 9))
    for i, idx in enumerate(idxs, 1):
        img = cv2.imread(paths[idx])[:, :, ::-1]
        true_lab = CLASS_NAMES[y_test[idx]]
        pred_lab = CLASS_NAMES[y_pred[idx]]
        plt.subplot(3, 4, i)
        plt.imshow(img)
        plt.title(f"True: {true_lab}\nPred: {pred_lab}",
                  color=("green" if true_lab == pred_lab else "red"))
        plt.axis('off')
    plt.tight_layout()
    plt.savefig(os.path.join(RESULT_DIR, "small_pred_samples.png"), dpi=300)
    plt.close()


if __name__ == "__main__":
    best_model = train()
    evaluate(best_model)
    print("全部完成！")